#include "kernel_operator.h"
#include <cmath>
//#include "math.h"

using namespace AscendC;

constexpr int32_t BUFFER_NUM = 2;      //昇腾双buffer技术

class KernelSoftmax {
public:
    __aicore__ inline KernelSoftmax() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y,int32_t size, int32_t x_ndarray[], int32_t x_dimensional, int x_dim)
    {

        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        //维度、axis、每个维度的个数、总元素个数
        this->x_dimensional = x_dimensional;
        this->dim = x_dim;
        for(int i = 0; i < x_dimensional; i++)
        {
            this->x_ndarray[i] = x_ndarray[i];
        }
        this->size =size;

        //维度转换
        if(this->dim < 0)
            this->dim += this->x_dimensional;

        xGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X *>(x), size + 8);
        yGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_Y *>(y), size + 8);

        int32_t cycles = 1;
        int32_t interval = 1;
        int32_t loopCount = 1;

        for(int32_t i = 0; i < this->x_dimensional; i++)      //如果是0，循环的就是第1维，1，1+2维，2，就是1+2+3维的总元素个数
        {
            loopCount *= this->x_ndarray[i];
        }
        loopCount = loopCount / this->x_ndarray[this->dim];   //循环的时候最外层的循环次数*最内层的循环

        cycles = this->x_ndarray[this->dim];                  //循环的时候并行层的循环数
        
        for(int32_t i = this->dim+1; i < this->x_dimensional; i++)
        {
            interval *= this->x_ndarray[i];     //跨其他维度时的跨度  //循环时需考虑元素间隔大小
        }

        //printf("cycle is %d, interval is %d, loopCount is %d\n",cycles, interval, loopCount);
        this->cycles = cycles;
        this->interval = interval;
        this->loopCount = loopCount;
        //printf("cycle is %d, interval is %d, loopCount is %d\n",this->cycles, this->interval, this->loopCount);

        //这里的分块逻辑是8192个x,1024个y和buff,也许不是最好，先用着
        //注意这里是interval为1，也就是连续才能并行，否则串行
        if(this->interval == 1)   
        {
            pipe.InitBuffer(inQueueX, BUFFER_NUM, 8192 * sizeof(DTYPE_X));
            pipe.InitBuffer(outQueueY, BUFFER_NUM, 8192 * sizeof(DTYPE_Y));
            pipe.InitBuffer(QueueBuff, 8192 * sizeof(DTYPE_X));
        }

    }

    // 计算阶乘
    __aicore__ float factorial(int n) {
        float result = 1.0f;
        for (int i = 1; i <= n; i++) {
            result *= i;
        }
        return result;
    }

    // 计算 e^r 的泰勒展开
    __aicore__ float exp_taylor(float r, int terms) {
        float result = 1.0f;
        float power = 1.0f;
        for (int i = 1; i < terms; i++) {
            power *= r;
            result += power / factorial(i);
        }
        return result;
    }

    // 精确计算 e^x
    __aicore__ float exp_precise(float x) {
        // 特殊值处理
        if (x > 88.0f) return HUGE_VALF;  // 溢出（float精度下的最大值约为 88）
        if (x < -103.0f) return 0.0f;    // 下溢（float精度下的最小值约为 -103）
        if (x == 0) return 1.0f;

        // 分解 x = n + r
        int n = (int)x;                 // 整数部分
        float r = x - n;                // 小数部分

        // 计算 e^n
        float exp_n = 1.0f;
        if (n > 0) {
            exp_n = 1.0f;
            for (int i = 0; i < n; i++) {
                exp_n *= 2.718281828459045f;  // 近似e
            }
        } else if (n < 0) {
            exp_n = 1.0f;
            for (int i = 0; i > n; i--) {
                exp_n /= 2.718281828459045f;  // 近似e
            }
        }

        // 计算 e^r
        float exp_r = exp_taylor(r, 10); // 使用 10 项泰勒展开

        // 返回结果
        return exp_n * exp_r;
    }



    __aicore__ inline void Process()
    {
        if constexpr (std::is_same_v<DTYPE_Y, float>)
        {
            //printf("is_same float\n");
            //这个设计就有点离谱，必须是8192能除尽的数才能并行,this->cycles不能为空

            //可并行部分  interval == 1 且 8192能除尽的数
            if((this->interval == 1) && (8192%this->cycles == 0) && (this->cycles <= 8192))
            {
                //printf("this->interval == 1\n");
                int capacity = 8192/this->cycles;    
                //printf("capacity is %d\n",capacity);
                for(int32_t i = 0; i < this->loopCount;)
                {
                    if((this->loopCount - i) >= capacity)
                    {
                        CopyIn(i, capacity);
                        Compute(i, capacity);
                        CopyOut(i, capacity);
                    }
                    else
                    {
                        CopyIn(i, this->loopCount - i);
                        Compute(i, this->loopCount - i);
                        CopyOut(i, this->loopCount - i);
                    }
                    i+=capacity;
                }
            }
            else
            {
                //printf("this->interval != 1\n");
                float temp_sum;
                float temp_add;
                float temp_max;
                for(int z = 0; z < this->loopCount; z++)
                {
                    int32_t x_num = z/this->interval;               //外层循环
                    x_num = x_num*this->cycles*this->interval + z%this->interval;  
                    //max
                    for(int i = 0; i < this->cycles; i++)
                    {
                        int32_t temp_index = x_num + i*this->interval;
                        temp_add = xGm.GetValue(temp_index);
                        
                        if(i == 0)
                        {
                            temp_max = temp_add;
                        }
                        else
                        {
                            if(temp_max < temp_add )
                                temp_max = temp_add;
                        }
                    }
                    //exp
                    for(int i = 0; i < this->cycles; i++)
                    {
                        int32_t temp_index = x_num + i*this->interval;
                        temp_add = xGm.GetValue(temp_index);
                        temp_add = temp_add - temp_max;
                        //temp_add = 1.0f + temp_add * (1.0f + temp_add * (0.5f + temp_add * (1.0f/6.0f + temp_add * (1.0f/24.0f + temp_add * 1.0f/120.0f))));
                        temp_add = exp_precise(temp_add);
                        xGm.SetValue(temp_index,temp_add);
                        
                        //xGm.SetValue(temp_index,(half)exp((float)temp_add-(float)temp_max));
                    }
                    //sum
                    for(int i = 0; i < this->cycles; i++)
                    {
                        int32_t temp_index = x_num + i*this->interval;
                        temp_add = xGm.GetValue(temp_index);
                        
                        if(i == 0)
                        {
                            temp_sum = temp_add;
                        }
                        else
                        {
                            temp_sum = temp_sum + temp_add;
                        }
                    }
                    //div
                    for(int i = 0; i < this->cycles; i++)
                    {
                        int32_t temp_index = x_num + i*this->interval;
                        temp_add = xGm.GetValue(temp_index);
                        yGm.SetValue(temp_index,(temp_add/temp_sum));
                    }
                    //yGm.SetValue(z, (DTYPE_Y)temp_sum);
                }
            }
        }
        else
        {
            //printf("is same half\n");
            float temp_sum;
            float temp_add;
            float temp_max;
            for(int z = 0; z < this->loopCount; z++)
            {
                int32_t x_num = z/this->interval;
                x_num = x_num*this->cycles*this->interval + z%this->interval;
                //max
                for(int i = 0; i < this->cycles; i++)
                {
                    int32_t temp_index = x_num + i*this->interval;
                    temp_add = xGm.GetValue(temp_index);
                    
                    if(i == 0)
                    {
                        temp_max = temp_add;
                    }
                    else
                    {
                        if(temp_max < temp_add )
                            temp_max = temp_add;
                    }
                }
                //printf("temp_max is %f\n",temp_max);
                //exp
                for(int i = 0; i < this->cycles; i++)
                {
                    int32_t temp_index = x_num + i*this->interval;
                    temp_add = xGm.GetValue(temp_index);
                    //printf("bef temp_add is %f\n",temp_add);
                    temp_add = temp_add - temp_max;
                    //temp_add = 1.0f + temp_add * (1.0f + temp_add * (0.5f + temp_add * (1.0f/6.0f + temp_add * (1.0f/24.0f + temp_add * 1.0f/120.0f))));
                    temp_add = exp_precise(temp_add);
                    xGm.SetValue(temp_index,(DTYPE_Y)temp_add);
                    //printf("exp temp_add is %f\n",temp_add);
                    
                    //xGm.SetValue(temp_index,(half)exp((float)temp_add-(float)temp_max));
                }
                //sum
                for(int i = 0; i < this->cycles; i++)
                {
                    int32_t temp_index = x_num + i*this->interval;
                    temp_add = xGm.GetValue(temp_index);
                    
                    if(i == 0)
                    {
                        temp_sum = temp_add;
                    }
                    else
                    {
                        temp_sum = temp_sum + temp_add;
                    }
                }
                //printf("temp_sum is %f\n",temp_sum);
                //div
                for(int i = 0; i < this->cycles; i++)
                {
                    int32_t temp_index = x_num + i*this->interval;
                    temp_add = xGm.GetValue(temp_index);
                    yGm.SetValue(temp_index,(DTYPE_Y) (temp_add/temp_sum));
                }
                //yGm.SetValue(z, (DTYPE_Y)temp_sum);
            }
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t i, int32_t capacity)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopy(xLocal, xGm[i*this->cycles], (this->cycles*capacity+7)/8*8);
        inQueueX.EnQue(xLocal);

    }

    __aicore__ inline void Compute(int32_t i, int32_t capacity)
    {
        LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();

        LocalTensor<float> xLocal = inQueueX.DeQue<float>();
        auto buff = QueueBuff.Get<float>();
        int cyl = this->cycles;


        float tmp = 0.0;
        for(int z=0;z<capacity;z++)
        {
            ReduceMax(yLocal[z], xLocal[this->cycles*z], buff, this->cycles); 
            tmp = yLocal[z].GetValue(0);
            //printf("max is %f\n",tmp);
            Duplicate(yLocal[this->cycles*z], tmp, this->cycles);
        }
        Sub(xLocal,xLocal,yLocal,8192);
        //printf("sub is %f\n",xLocal(0));
        Exp(xLocal,xLocal,8192);
        //printf("exp is %f\n",xLocal(0));
        for(int z=0;z<capacity;z++)
        {
            ReduceSum(yLocal[z], xLocal[this->cycles*z], buff, this->cycles); 
            tmp = yLocal[z].GetValue(0);
            //printf("sum is %f\n",tmp);
            Duplicate(yLocal[this->cycles*z], tmp, this->cycles);
        }
        Div(yLocal,xLocal,yLocal,8192);
        //printf("div is %f\n",yLocal(0));
        
        outQueueY.EnQue<float>(yLocal);
        inQueueX.FreeTensor(xLocal);
    }
    
    __aicore__ inline void CopyOut(int32_t i, int32_t j)
    {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        //DataCopy(yGm[i], yLocal, (j+7)/8*8);
        DataCopy(yGm[i*this->cycles], yLocal, (this->cycles*j+7)/8*8);
        //printf("yGm %d is %f\n",i, yGm(i));
        outQueueY.FreeTensor(yLocal);
    }


private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueBuff;
    
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;

    int32_t size;
    int32_t x_ndarray[10];
    int32_t x_dimensional;
    int dim;

    int32_t cycles;
    int32_t interval;
    int32_t loopCount;
};

extern "C" __global__ __aicore__ void softmax(GM_ADDR x, GM_ADDR y,  GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelSoftmax op;
    //补充init和process函数调用内容
    op.Init(x, y,tiling_data.size, tiling_data.x_ndarray, tiling_data.x_dimensional, tiling_data.x_dim);
    op.Process();
}

